import ImageProcess
import sys
import PyQt5.Qt as Qt
from PyQt5.QtGui import *
from PyQt5.QtCore import *
from PyQt5.QtWidgets import *

class Canvas(QWidget):
    def __init__(self, parent):
        super(Canvas, self).__init__()
        self.setFixedSize(560, 560)
        self.setParent(parent)
        self.onDrawing = False
        self.setGeometry(10, 10, 560, 560)
        self.setStyleSheet("background-color:rgb(0,0,0);")
        self.pen = QPen()
        self.pen.setColor(QColor(255, 255, 255))
        self.pen.setWidth(30)
        self.pen.setJoinStyle(Qt.RoundJoin)
        self.pen.setCapStyle(Qt.RoundCap)
        self.buffer = QPixmap(560, 560)
        self.painter = QPainter()
        self.temp = QPainter()
        self.prev = QPoint()
        self.current = QPoint()

    def paintEvent(self, event):
        if self.onDrawing:
            self.temp.drawLine(self.prev, self.current)
        self.painter.begin(self)
        self.painter.setRenderHint(QPainter.Antialiasing)
        self.painter.setPen(self.pen)
        self.painter.drawPixmap(self.rect(), self.buffer)
        self.painter.end()
        return super().paintEvent(event)

    def mouseMoveEvent(self, event: QMouseEvent):
        if self.onDrawing:
            self.current = event.pos()
            self.repaint()
            self.prev = self.current
        return super().mouseMoveEvent(event)

    def mousePressEvent(self, event: QMouseEvent):
        if event.button() == Qt.LeftButton:
            self.temp.begin(self.buffer)
            self.temp.setRenderHint(QPainter.Antialiasing)
            self.temp.setPen(self.pen)
            self.onDrawing = True
            self.prev = self.current = event.pos()
            self.update()
        return super().mousePressEvent(event)

    def mouseReleaseEvent(self, event: QMouseEvent):
        if event.button() == Qt.LeftButton:
            self.update()
            self.temp.end()
            self.onDrawing = False
        return super().mouseReleaseEvent(event)

    def setWidth(self, w):
        self.pen.setWidth(w)

    def clear(self):
        self.buffer.fill(QColor(0, 0, 0))
        self.update()


class MainWindow(QMainWindow):
    def __init__(self):
        super(MainWindow, self).__init__()
        self.setWindowTitle("MNIST-GUI")
        self.setFixedSize(690, 580)
        self.canvas = Canvas(self)
        self.btn = QPushButton(self)
        self.btn.setGeometry(580, 10, 100, 50)
        self.btn.setText("Detect")
        self.clearBtn = QPushButton(self)
        self.clearBtn.setGeometry(580, 70, 100, 50)
        self.clearBtn.setText("Clear")
        self.label = QLabel(self)
        self.label.setGeometry(580, 130, 100, 50)
        self.label.setText("None")
        self.label.setStyleSheet("background-color: white;")
        self.label.setAlignment(Qt.AlignCenter)
        self.slider = QSlider(self)
        self.slider.setGeometry(580, 200, 100, 30)
        self.slider.setMinimum(1)
        self.slider.setMaximum(100)
        self.slider.setValue(30)
        self.slider.setOrientation(Qt.Horizontal)
        # self.btn.clicked.connect(lambda: self.canvas.buffer.toImage().save('img.jpg'))
        self.btn.clicked.connect(lambda: self.label.setText(ImageProcess.predict(self.canvas.buffer)))
        self.clearBtn.clicked.connect(lambda: self.canvas.clear())
        self.slider.valueChanged.connect(
            lambda: self.canvas.setWidth(self.slider.value()))
        